from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging

import copy
import torch
import torch.nn.functional as F
from torch import nn, Tensor
from typing import Optional

import math
from models.pose.intraNet import build_intraNet

from utils.utils import get_valid_output
from utils import ddp_opx


BN_MOMENTUM = 0.1
logger = logging.getLogger(__name__)


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")


class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight=True):
        super(JointsMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='mean')
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        batch_size = output.size(0)
        num_joints = output.size(1)
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
        loss = 0

        for idx in range(num_joints):
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss += 0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                )
            else:
                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

        return loss / num_joints


class DeConv(nn.Module):
    def __init__(self, cfg): 
        super().__init__() 
        extra = cfg.MODEL.EXTRA
        self.deconv_with_bias = extra.DECONV_WITH_BIAS

        self.heatmap_size = cfg.MODEL.HEATMAP_SIZE
        self.trans_size = cfg.MODEL.TRANS_SIZE
        
        mod = self.heatmap_size[0] // self.trans_size[0]
        self.layer_num = int(math.log(mod, 2))
        
        self.deconv_layers = nn.ModuleList(self._make_deconv_layer(
            extra.NUM_DECONV_LAYERS,   # 1
            extra.NUM_DECONV_FILTERS,  # [d_model]
            extra.NUM_DECONV_KERNELS,  # [4]
        ) for _ in range(self.layer_num))
        
    def _get_deconv_cfg(self, deconv_kernel, index):
        if deconv_kernel == 4:
            padding = 1
            output_padding = 0
        elif deconv_kernel == 3:
            padding = 1
            output_padding = 1
        elif deconv_kernel == 2:
            padding = 0
            output_padding = 0

        return deconv_kernel, padding, output_padding
        
    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
        assert num_layers == len(num_filters), \
            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
        assert num_layers == len(num_kernels), \
            'ERROR: num_deconv_layers is different len(num_deconv_filters)'

        layers = []
        for i in range(num_layers):
            kernel, padding, output_padding = \
                self._get_deconv_cfg(num_kernels[i], i)

            planes = num_filters[i]
            layers.append(
                nn.ConvTranspose2d(
                    in_channels=planes,
                    out_channels=planes,
                    kernel_size=kernel,
                    stride=2,
                    padding=padding,
                    output_padding=output_padding,
                    bias=self.deconv_with_bias))
            layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
            layers.append(nn.ReLU(inplace=True))
            self.inplanes = planes

        return nn.Sequential(*layers)
    
    def forward(self, x):
        for layer in self.deconv_layers:
            x = layer(x)
        return x


class TransformerEncoderLayer(nn.Module):
    """ Modified from https://github.com/facebookresearch/detr/blob/master/models/transformer.py"""

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False, return_atten_map=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before
        self.return_atten_map = return_atten_map

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(src, pos)
        if self.return_atten_map:
            src2, att_map = self.self_attn(q, k, value=src,
                                           attn_mask=src_mask,
                                           key_padding_mask=src_key_padding_mask)
        else:
            src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                                  key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        if self.return_atten_map:
            return src, att_map
        else:
            return src

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        if self.return_atten_map:
            src2, att_map = self.self_attn(q, k, value=src,
                                           attn_mask=src_mask,
                                           key_padding_mask=src_key_padding_mask)
        else:
            src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                                  key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        if self.return_atten_map:
            return src, att_map
        else:
            return src

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)


class TransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers,
                 norm=None, pe_only_at_begin=False, return_atten_map=False):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.pe_only_at_begin = pe_only_at_begin
        self.return_atten_map = return_atten_map
        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
                
    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        output = src
        atten_maps_list = []
        for layer in self.layers:
            if self.return_atten_map:
                output, att_map = layer(output, src_mask=mask, pos=pos,
                                        src_key_padding_mask=src_key_padding_mask)
                atten_maps_list.append(att_map)
            else:
                output = layer(output, src_mask=mask,  pos=pos,
                               src_key_padding_mask=src_key_padding_mask)

            # only add position embedding to the first atttention layer
            pos = None if self.pe_only_at_begin else pos

        if self.norm is not None:
            output = self.norm(output)

        if self.return_atten_map:
            return output, torch.stack(atten_maps_list)
        else:
            return output


class InterNet(nn.Module):

    def __init__(self, cfg, **kwargs):
        super(InterNet, self).__init__()
        extra = cfg.MODEL.EXTRA
        self.trans_size = cfg.MODEL.TRANS_SIZE
        self.heatmap_size = cfg.MODEL.HEATMAP_SIZE
        d_model = cfg.MODEL.DIM_MODEL
        inter_layers = cfg.MODEL.ENCODER_MULTI_LAYERS

        self.intraNet = build_intraNet(cfg)    # HRNet(*)、TransPos、HRFormer
        
        encoder_layer = TransformerEncoderLayer(d_model=d_model, nhead=cfg.MODEL.N_HEAD, 
                        dim_feedforward=cfg.MODEL.DIM_FEEDFORWARD, activation='relu')
        self.interNet = TransformerEncoder(encoder_layer, inter_layers)
    
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        self.upsample_layer = DeConv(cfg)

        self.final_layer = nn.Conv2d(
            in_channels=d_model,
            out_channels=cfg['MODEL']['NUM_JOINTS'],
            kernel_size=extra['FINAL_CONV_KERNEL'],
            stride=1,
            padding=1 if extra['FINAL_CONV_KERNEL'] == 3 else 0
        )

        self.criterion = JointsMSELoss(use_target_weight=True)

    def max_pool(self, x, mod,):
        rate = int(math.log(mod, 2))
        for _ in range(rate):
            x = self.maxpool(x)
        return x

    def padding_tensor(self, tensor, length):
        device = tensor.device
        N = max(length)
        res_list = []
        mask_list = []
        tensor_shape = tuple(tensor.shape[-3:])
        mask_shape = tuple(tensor.shape[-2:])

        tensor_split = tensor.split(length, dim=0)

        for valid_item in tensor_split:
            n = len(valid_item)
            if n == 0:
                invalid = [torch.ones(mask_shape, dtype=torch.bool) for _ in range(N)]
                mask_item = torch.stack(invalid, dim=0)

                padding_list = [torch.zeros(tensor_shape, dtype=torch.float32) for _ in range(N)]
                padding_list = torch.stack(padding_list, dim=0)
                padding_list = padding_list.to(device)
                valid_item = torch.cat((valid_item, padding_list), dim=0)
            elif n == N:
                valid = [torch.zeros(mask_shape, dtype=torch.bool) for _ in range(n)]
                mask_item = torch.stack(valid, dim=0)
            else:   # need padding
                padding_list = [torch.zeros(tensor_shape, dtype=torch.float32) for _ in range(N-n)]
                padding_list = torch.stack(padding_list, dim=0)
                padding_list = padding_list.to(device)
                valid_item = torch.cat((valid_item, padding_list), dim=0)

                valid = [torch.zeros(mask_shape, dtype=torch.bool) for _ in range(n)]
                mask_item = torch.stack(valid, dim=0)
                invalid = [torch.ones(mask_shape, dtype=torch.bool) for _ in range(N-n)]
                invalid = torch.stack(invalid, dim=0)
                mask_item = torch.cat((mask_item, invalid), dim=0)
            res_list.append(valid_item)
            mask_list.append(mask_item)

        res = torch.stack(res_list, dim=0)
        mask = torch.stack(mask_list, dim=0)
        return res, mask.to(device)

    def forward(self, samples):
        if isinstance (samples, dict):
            return self.forward_parallel(samples)
        elif isinstance (samples, list):
            return self.forward_list(samples)

    def forward_list(self, x_list):
        keypoints_loss = 0
        for idx, item in enumerate(x_list):
            if item is None:
                x_list[idx] = None
                continue
            x = item.pop('input')
            # ==> intra
            x, output_single = self.intraNet(x)
            # outputs['single'].append(output_single)
            x_single = x

            # ==> inter
            x = self.max_pool(x, x.shape[-1] // self.trans_size[0])
            bs, c, h, w = x.shape               # L = h*w
            x = x.flatten(2).permute(2, 0, 1)   # [L, bs, c]
            x = self.interNet(x)
            x = x.permute(1, 2, 0).contiguous().view(bs, c, h, w)
            x_list[idx]['feature'] = x

            x_multi = self.upsample_layer(x) + x_single      # inter + intra

            x_multi_out = self.final_layer(x_multi)
            if 'target' in item:
                keypoints_loss += self.criterion(x_multi_out, item['target'], item['target_weight'])
                item.pop('target')
                item.pop('target_weight')

        # outputs['multi'] = torch.cat(outputs['multi'], dim=0)
        # outputs['single'] = torch.cat(outputs['single'], dim=0)
        return x_list, keypoints_loss

    def forward_parallel(self, samples):
        outputs = {'single': None, 'multi': None}
        length = samples['length']
        x = samples['x']

        x, outputs['single'] = self.intraNet(x)
        single_res = x

        x = self.max_pool(x, x.shape[-1] // self.trans_size[-2])        # down  [64,48] -> [16,12]
        x_padding, mask = self.padding_tensor(x, length) # [bs, n, c, h, w]

        b, n, c, h, w = x_padding.shape
        x_padding = x_padding.permute(0, 2, 1, 3, 4).flatten(2).permute(2, 0, 1)    # [n*h*w, B, C]
        mask = mask.flatten(1) if mask is not None else None
        x_padding = self.interNet(x_padding, src_key_padding_mask=mask, pos=None)    # [L, bs, c]
        x_padding = x_padding.permute(1, 2, 0).contiguous().view(b, c, n, h, w)
        x_padding = x_padding.permute(0, 2, 1, 3, 4)

        x_valid = get_valid_output(x_padding, length)
        x_valid = self.upsample_layer(x_valid)                          # up    [16,12] -> [64,48]

        x_valid = single_res + x_valid      # Residual # B,96,64,48
        
        outputs['multi'] = self.final_layer(x_valid)

        return outputs


def get_pose_net(cfg, **kwargs):

    model = InterNet(cfg, **kwargs)
    
    return model

def build_pose_net(cfg, **kwargs):
    model = InterNet(cfg, **kwargs)
    checkpoint_file = 'checkpoints/pose_hrnet_inter2.pth'
    if ddp_opx.is_main_process():
        logger.info(f'=> pose net load weights from {checkpoint_file}')
    ckpt_state_dict = torch.load(checkpoint_file, map_location=torch.device('cpu'))
    model.load_state_dict(ckpt_state_dict, strict=True)
    return model
